package org.example.learn.javaagent.javassist;

import javassist.*;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;


// ClassFileTransformer 接口作用： 代理提供此接口的实现以转换类文件。 转换发生在 JVM 定义类之前。
// 增加方法字节 实现监控方法耗时
public class MonitorTransformer implements ClassFileTransformer {

    private static final Set<String> classNameSet = new HashSet<String>();

    static {
        classNameSet.add("org.example.learn.javaagent.javassist.Run");
    }

    // 转换类文件处理方法 接口实现方法
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        String currentClassName = className.replaceAll("/",".");
        if (!classNameSet.contains(currentClassName)){ // 检查classNameSet中是否含有该类
            return null;
        }
        System.out.println("transform: [" + currentClassName + "]");
        try {
            // 获取Class
            CtClass ctClass = ClassPool.getDefault().get(currentClassName);
            // CtBehavior表示方法、构造函数或静态构造函数（类初始值设定项）。 它是CtMethod和CtConstructor的抽象超类。要直接读取或修改字节码，请获取MethodInfo对象
            // 获取该类的 方法信息（方法、构造方法、静态构造方法） 数组
            CtBehavior[] ctBehavior = ctClass.getDeclaredBehaviors();
            for (CtBehavior behavior:ctBehavior) {
                enhanceMethod(behavior);
            }
            // Class 转换为 byte[]
            return ctClass.toBytecode();
        } catch (NotFoundException e) {
            e.printStackTrace();
        } catch (CannotCompileException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    // 处理方法 添加增加字节信息
    private void enhanceMethod(CtBehavior behavior) throws CannotCompileException {
        if (behavior.isEmpty()){
            return;
        }
        String methodName = behavior.getName();
        if ("main".equals(methodName)){
            return;
        }
        final StringBuilder source = new StringBuilder();
        source.append("{");
        // 前置增强：插入时间戳
        source.append("long startTime = System.nanoTime();");
        // 调用原有方法，类似于method(); $$ 表示所有参数
        source.append("$_=$proceed($$);\n");
        source.append("System.out.print(\"method:[");
        source.append(methodName);
        source.append("]\");");
        source.append("\n");
        // 计算耗时
        source.append("System.out.println(\" const:[\" + (System.nanoTime() - startTime) + \"ns]\");");
        source.append("}");
        ExprEditor editor = new ExprEditor() {
            @Override
            public void edit(MethodCall m) throws CannotCompileException {
                m.replace(source.toString());
            }
        };
        behavior.instrument(editor);
    }
}
